### 第六章 PyTorch 进阶训练技巧 ###
# 6.1 自定义损失函数 #

# 6.1.1 以函数方式定义 #

def my_loss(output, target):
    loss = torch.mean((output - target)**2)
    return loss

# 6.1.2 以类方式定义 #

# 虽然以函数定义的方式很简单，但是以类方式定义更加常用，
# 在以类方式定义损失函数时，
# 我们如果看每一个损失函数的继承关系我们就可以发现Loss函数部分继承自_loss, 部分继承自_WeightedLoss,
# 而_WeightedLoss继承自_loss， _loss继承自 nn.Module。
# 我们可以将其当作神经网络的一层来对待，同样地，我们的损失函数类就需要继承自nn.Module类

###  $$ DSC = \frac{2|X∩Y|}{|X|+|Y|} $$
class DiceLoss(nn.Module):
    def __init__(self, weight = None, size_average = Ture):
        super(DiceLoss, self).__init__()

        def forward(self, inputs, targets, smooth = 1):
            inputs = F.sigmoid(inputs)
            inputs = inputs.view(-1)
            targets = targets.view(-1)
            intersection = (inputs * targets).sum()
            dice = (2.*intersection + smooth)/(inputs.sum() + targets.sum() + smooth)
            return 1 - dice

# 使用方法
criterion = DiceLoss()
loss = criterion(inputs, targets)

# 除此之外，常见的损失函数还有BCE-Dice Loss，Jaccard/Intersection over Union (IoU) Loss，Focal Loss......

class DiceBCELoss(nn.Module):
    def __init__(self, weight = None, size_average = True):
        super(DiceBCELoss, self).__init__()

    def forward(self, inouts, targets, smooth = 1):
        inputs = F.sigmoid(inputs)
        inputs = inputs.view(-1)
        targets = targets.view(-1)
        intersection = (inputs * targets).sum()
        dice_loss = 1 - (2. * intersection + smmoth)/(inputs.sum() + targets.sum() + smooth)

class IoULoss(nn.Module):
    def __init__(self, weight = None, size_average = True):
        super(IoULoss, self).__init__()

    def forward(self, inputs, targets, smooth = 1):
        inputs = F.sigmoid(inputs)
        inputs = inputs.view(-1)
        targets = targets.view(-1)
        intersection = (inputs * targets).sum()
        total = (inputs + targets).sum()
        union = total - intersection

        Iou = (intersection + smooth)/(union + smooth)

        return  1 - IoU

ALPHA = 0.8
GAMMA = 2

class FocalLoss(nn.Module):
    def __init__(self, weight = None, size_average = True):
        super(FocalLoss, self).__init__()

    def forward(self, inputs, targets, alpha = ALPHA, gamma = GAMMA, smooth = 1)
        inputs = F.sigmoid(inputs)
        inputs = inputs.view(-1)
        targets = targets.view(-1)
        BCE = F.binary_cross_entropy(inputs, targets, reduction='mean')
        BCE_EXP = torch.exp(-BCE)
        focal_loss = alpha * (1-BCE_EXP)**gamma * BCE

        return focal_loss

# 在自定义损失函数时，涉及到数学运算时，我们最好全程使用PyTorch提供的张量计算接口，
# 这样就不需要我们实现自动求导功能并且我们可以直接调用cuda，
# 使用numpy或者scipy的数学运算时，操作会有些麻烦，


""""
本节参考
1. https://www.kaggle.com/bigironsphere/loss-function-library-keras-pytorch/notebook
2. https://www.zhihu.com/question/66988664/answer/247952270
3. https://blog.csdn.net/dss_dssssd/article/details/84103834
4. https://zj-image-processing.readthedocs.io/zh_CN/latest/pytorch/%E8%87%AA%E5%AE%9A%E4%B9%89%E6%8D%9F%E5%A4%B1%E5%87%BD%E6%95%B0/
5. https://blog.csdn.net/qq_27825451/article/details/95165265
6. https://discuss.pytorch.org/t/should-i-define-my-custom-loss-function-as-a-class/89468
""""

# 6.2 动态调整学习率 #

# 学习率的选择是深度学习中一个困扰人们许久的问题，学习速率设置过小，会极大降低收敛速度，增加训练时间；
# 学习率太大，可能导致参数在最优解两侧来回振荡。
# 但是当我们选定了一个合适的学习率后，经过许多轮的训练后，可能会出现准确率震荡或loss不再下降等情况，
# 说明当前学习率已不能满足模型调优的需求。
# 此时我们就可以通过一个适当的学习率衰减策略来改善这种现象，提高我们的精度。
# 这种设置方式在PyTorch中被称为scheduler，也是我们本节所研究的对象。


# 6.2.1 使用官方scheduler


""""
在训练神经网络的过程中，学习率是最重要的超参数之一，作为当前较为流行的深度学习框架，
PyTorch已经在torch.optim.lr_scheduler为我们封装好了一些动态调整学习率的方法供我们使用，
如下面列出的这些scheduler。

lr_scheduler.LambdaLR
lr_scheduler.MultiplicativeLR
lr_scheduler.StepLR
lr_scheduler.MultiStepLR
lr_scheduler.ExponentialLR
lr_scheduler.CosineAnnealingLR
lr_scheduler.ReduceLROnPlateau
lr_scheduler.CyclicLR
lr_scheduler.OneCycleLR
lr_scheduler.CosineAnnealingWarmRestarts
""""

# 选择一种优化器
optimizer = torch.optim.Adam(...)
# 选择一种上述的一种或多种动态调整学习率的方法
scheduler1 = torch.optim.lr_scheduler....
scheduler2 = torch.optim.lr_scheduler....
...
schedulern = torch.optim.lr_scheduler....
# 进行训练
for epoch in range(100):
    train(...)
    validate(...)
    optimizer.step()
    # 需要在优化器参数更新后再动态调整学习率
    scheduler1.step()
    ...
    schedulern.step()

# 我们在使用官方给出的torch.optim.lr_scheduler时，
# 需要将scheduler.step()放在optimizer.step()后面进行使用。

# 6.2.2 自定义scheduler

# 虽然PyTorch官方给我们提供了许多的API，
# 但是在实验中也有可能碰到需要我们自己定义学习率调整策略的情况，
# 而我们的方法是自定义函数adjust_learning_rate来改变param_group中lr的值，
# 在下面的叙述中会给出一个简单的实现。

# 假设我们现在正在做实验，需要学习率每30轮下降为原来的1/10，
# 假设已有的官方API中没有符合我们需求的，那就需要自定义函数来实现学习率的改变。

def adjust_learning_rate(optimizer, epoch):
    lr = args.lr * (0.1 ** (epoch // 30))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

# 有了adjust_learning_rate函数的定义，在训练的过程就可以调用我们的函数来实现学习率的动态变化

def adjust_learning_rate(optimizer,...):
    ...
optimizer = torch.optim.SGD(model.parameters(), lr = args.lr, momentum = 0.9)
for epoch in range(10):
    train(...)
    validate(...)
    adjust_learning_rate(optimizer, epoch)

#  本节参考内容：
#  PyTorch官方文档https://pytorch.org/docs/stable/optim.html


# 6.3 模型微调 #

"""
深度学习的发展导致模型的参数越来越多，许多开源模型都是在较大数据集上进行训练的
但在实际应用中，我们的数据集可能只有几千张，
这时从头开始训练具有几千万参数的大型神经网络是不现实的，
因为越大的模型对数据量的要求越大，过拟合无法避免。

应用迁移学习，将从源数据集学到的知识迁移到目标数据集上，或许是一个好主意
例如，虽然ImageNet数据集的图像大多跟椅子无关，
但在该数据集上训练的模型可以抽取较通用的图像特征，
从而能够帮助识别边缘、纹理、形状和物体组成等。
这些类似的特征对于识别椅子也可能同样有效。

迁移学习的一大应用场景是模型微调（finetune）。
简单来说，就是我们先找到一个同类的别人训练好的模型，把别人现成的训练好了的模型拿过来，
换成自己的数据，通过训练调整一下参数。
在PyTorch中提供了许多预训练好的网络模型（VGG，ResNet系列，mobilenet系列......），
这些模型都是PyTorch官方在相应的大型数据集训练好的。
学习如何进行模型微调，可以方便我们快速使用预训练模型完成自己的任务。
"""

# 6.3.1 模型微调的流程

"""
1.在源数据集(如ImageNet数据集)上预训练一个神经网络模型，即源模型。
2.创建一个新的神经网络模型，即目标模型。它复制了源模型上除了输出层外的所有模型设计及其参数。
我们假设这些模型参数包含了源数据集上学习到的知识，且这些知识同样适用于目标数据集。
我们还假设源模型的输出层跟源数据集的标签紧密相关，因此在目标模型中不予采用。
3.为目标模型添加一个输出⼤小为⽬标数据集类别个数的输出层，并随机初始化该层的模型参数。
4.在目标数据集上训练目标模型。我们将从头训练输出层，而其余层的参数都是基于源模型的参数微调得到的。
"""

# 6.3.2 使用已有的模型结构

# 实例化网络

import torchvision.models as models
resnet18 = models.resnet18()
# resnet18 = models.resnet18(pretrained = False) 等价于与上面的表达式
alexnet = models.alexnet()
vgg16 = models.vgg16()
squeezenet = models.squeezenet1_0()
densenet = models.densenet161()
inception = models.inception_v3()
googlenet = models.googlenet()
shufflenet = models.shufflenet_v2_x1_0()
mobilenet_v2 = models.mobilenet_v2()
mobilenet_v3_large = models.mobilenet_v3_large()
mobilenet_v3_small = models.mobilenet_v3_small()
resnext50_32x4d = models.resnext50_32x4d()
wide_resnet50_2 = models.wide_resnet50_2()
mnasnet = models.mnasnet1_0()

# 传递预训练参数（pretrained）

# 通过True或者False来决定是否使用预训练好的权重，
# 在默认状态下pretrained = False，意味着我们不使用预训练得到的权重，
# 当pretrained = True，意味着我们将使用在一些数据集上预训练得到的权重。

import trochvision.models as models
resnet18 = models.resnet18(pretrained=True)
alexnet = models.alexnet(pretrained=True)
squeezenet = models.squeezenet1_0(pretrained=True)
vgg16 = models.vgg16(pretrained=True)
densenet = models.densenet161(pretrained=True)
inception = models.inception_v3(pretrained=True)
googlenet = models.googlenet(pretrained=True)
shufflenet = models.shufflenet_v2_x1_0(pretrained=True)
mobilenet_v2 = models.mobilenet_v2(pretrained=True)
mobilenet_v3_large = models.mobilenet_v3_large(pretrained=True)
mobilenet_v3_small = models.mobilenet_v3_small(pretrained=True)
resnext50_32x4d = models.resnext50_32x4d(pretrained=True)
wide_resnet50_2 = models.wide_resnet50_2(pretrained=True)
mnasnet = models.mnasnet1_0(pretrained=True)

"""
注意事项：

1. 通常PyTorch模型的扩展为.pt或.pth，程序运行时会首先检查默认路径中是否有已经下载的模型权重，
一旦权重被下载，下次加载就不需要下载了。

2. 一般情况下预训练模型的下载会比较慢，我们可以直接通过迅雷或者其他方式去
https://github.com/pytorch/vision/tree/master/torchvision/models
 查看自己的模型里面model_urls，然后手动下载，
 预训练模型的权重在Linux和Mac的默认下载路径是用户根目录下的.cache文件夹。
 在Windows下就是C:\Users\<username>\.cache\torch\hub\checkpoint。
 我们可以通过使用 torch.utils.model_zoo.load_url()设置权重的下载地址。

3.如果觉得麻烦，还可以将自己的权重下载下来放到同文件夹下，然后再将参数加载网络。

self.model = models.resnet50(pretrained=False)
self.model.load_state_dict(torch.load('./model/resnet50-19c8e357.pth'))

4.如果中途强行停止下载的话，一定要去对应路径下将权重文件删除干净，要不然可能会报错。
"""

# 6.3.3 训练特定层

"""
在默认情况下，参数的属性.requires_grad = True，如果我们从头开始训练或微调不需要注意这里。
但如果我们正在提取特征并且只想为新初始化的层计算梯度，其他参数不进行改变。
那我们就需要通过设置requires_grad = False来冻结部分层。
在PyTorch官方中提供了这样一个例程。
"""

def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False #此例程中模型参数未冻结

"""
在下面我们仍旧使用resnet18为例的将1000类改为4类，
但是仅改变最后一层的模型参数，不改变特征提取的模型参数；
注意我们先冻结模型参数的梯度，再对模型输出部分的全连接层进行修改，
这样修改后的全连接层的参数就是可计算梯度的。
"""

import torchvision.models as models
# 冻结参数的梯度
feature_extracting = True
model = models.resnet18(pretrained = True)
set_parameter_requires_grad(model, feature_extracting)
# 修改模型
num_ftrs = model.fc.in_features
model.fc = nn.Linear(in_features = 512, out_features = 4, bias = True)

"""
之后在训练过程中，model仍会进行梯度回传，但是参数更新则只会发生在fc层。
通过设定参数的requires_grad属性，我们完成了指定训练模型的特定层的目标，
这对实现模型微调非常重要。
"""

"""
本节参考
1、参数更新：
torch.optim - PyTorch 中文文档
https://www.pytorchtutorial.com/docs/package_references/torch-optim/

2、给不同层分配不同的学习率
[ pytorch ] 基本使用丨8. 优化器optimizer的使用丨_小小的行者的博客-CSDN博客
https://blog.csdn.net/jdzwanghao/article/details/90402577

"""


# 6.4 半精度训练 #

"""

我们提到PyTorch时候，总会想到要用硬件设备GPU的支持，也就是“卡”。
GPU的性能主要分为两部分：算力和显存，前者决定了显卡计算的速度，后者则决定了显卡可以同时放入多少数据用于计算。

在可以使用的显存数量一定的情况下，每次训练能够加载的数据更多（也就是batch size更大），则也可以提高训练效率。
另外，有时候数据本身也比较大（比如3D图像、视频等），显存较小的情况下可能甚至batch size为1的情况都无法实现。
因此，合理使用显存也就显得十分重要。

我们观察PyTorch默认的浮点数存储方式用的是torch.float32，小数点后位数更多固然能保证数据的精确性，
但绝大多数场景其实并不需要这么精确，只保留一半的信息也不会影响结果，也就是使用torch.float16格式。
由于数位减了一半，因此被称为“半精度”.
"""
# import autocast
from torch.cuda.amp import autocast
# 模型设置
# 在模型定义中，使用python的装饰器方法，用autocast装饰模型中的forward函数
@autocast
def forward(self, x):
    ...
    return x
# 训练过程
#在训练过程中，只需在将数据输入模型及其之后的部分放入“with autocast():“即可：
for x in train_loader:
    x = x.cuda()
    with autocast():
    output = model(x)
    ...

"""
注意：

半精度训练主要适用于数据本身的size比较大（比如说3D图像、视频等）。
当数据本身的size并不大时（比如手写数字MNIST数据集的图片尺寸只有28*28），
使用半精度训练则可能不会带来显著的提升。

"""
